import random
import numpy as np
import pandas as pd
from General.DatabaseReader import DatabaseReader

np.set_printoptions(threshold=np.inf)


# gap是分割时间的间隔，转换时间到时间戳
def convert_to_timestamp(time, gap):
    minutes = time.hour * 60 + time.minute
    return minutes // gap


class Learning:
    """
        该强化学习订单分配算法分为两个步骤，learning和planning
        1）learning：
        学习阶段，目的是学习机器人的状态价值，即V_s，V_s是一个二维表格，V_s[i][j]表示机器人当前所处的时间段i、地区j的状态价值
        地区是对点位进行聚类之后得到的一个抽象地理表示，以减少表格的规模，在Clustering.NodeClusterHandler中完成聚类，结果保存在data.node_index_with_cluster中
    """
    def __init__(self, time_count, db_reader, code_index_map_filepath, weight_params):
        """
        初始化 Learning 类
        :param time_count: 一天中的时间步数
        :param db_reader: 数据库读取器实例
        :param code_index_map_filepath: code_index_map 文件路径
        :param weight_params: 用于加权平均的权重参数，格式为 {'weight_1': 0.4, 'weight_2': 0.3, 'weight_3': 0.3}
        """
        self.time_count = time_count  # 一天的时间步数
        self.db_reader = db_reader  # 数据库读取类的实例
        self.code_index_map_filepath = code_index_map_filepath  # code_index_map 文件路径
        self.weight_params = weight_params  # 加权平均参数

        # 从 code_index_map 文件中读取集群数据，计算区域数
        code_index_map = pd.read_excel(self.code_index_map_filepath)
        self.region_count = code_index_map['cluster'].nunique()  # 计算唯一的 cluster 数量作为区域数

        # 从数据库读取 robot_order 数据
        self.df = self.db_reader.read_robot_order()
        # 小规模数据，测试用
        # self.df = self.db_reader.read_robot_order_mini()



        # 确保时间字段类型一致，将所有时间列转换为 datetime 类型
        self.df['create_time'] = pd.to_datetime(self.df['create_time'], errors='coerce')
        self.df['start_time'] = pd.to_datetime(self.df['start_time'], errors='coerce')
        self.df['put_time'] = pd.to_datetime(self.df['put_time'], errors='coerce')
        self.df['get_time'] = pd.to_datetime(self.df['get_time'], errors='coerce')
        self.df['finish_time'] = pd.to_datetime(self.df['finish_time'], errors='coerce')

        self.orders = []  # 订单列表
        if 24 * 60 % time_count != 0:
            print('请选择合适的时间总步长')
        self.time_gap = 24 * 60 / time_count  # 计算每个时间步长的时间间隔

        self.D_backup = []
        self.N = np.ones([self.time_count, self.region_count])  # 记录次数矩阵
        self.V_s = np.zeros([self.time_count, self.region_count])  # 价值矩阵，表示(时间，地区）的价值大小

    def get_order_data(self):
        """
        获取订单数据，使用 code_index_map 进行映射，并计算订单的价值
        :return: 格式化后的订单列表
        """
        file_a = self.df
        # 从 Excel 文件中读取 code_index_map，用于映射 start_node 和 end_node 到区域
        file_b = pd.read_excel(self.code_index_map_filepath)

        # 将集群数据转换为字典，进行区域映射
        cluster_dict = file_b.set_index('code')['cluster'].to_dict()

        # 将 start_node 和 end_node 映射为 start_region 和 end_region
        file_a['start_region'] = file_a['start_node'].map(cluster_dict)
        file_a['end_region'] = file_a['end_node'].map(cluster_dict)

        # 删除那些在 cluster_dict 中没有匹配到 region 的行（即 start_region 或 end_region 为 NaN 的行）
        file_a = file_a.dropna(subset=['start_region', 'end_region'])

        # 将 start_time 和 finish_time 转换为时间戳
        file_a['start_timestamp'] = file_a['create_time'].apply(lambda x: convert_to_timestamp(x, self.time_gap))
        file_a['end_timestamp'] = file_a['finish_time'].apply(lambda x: convert_to_timestamp(x, self.time_gap))

        # 初始化 used 字段为 0
        file_a['used'] = 0

        # 计算订单价值/奖励，使用传入的权重参数
        def calculate_order_value(row):
            time_diff_1 = (row['start_time'] - row['create_time']).total_seconds() / 60.0
            time_diff_2 = (row['put_time'] - row['start_time']).total_seconds() / 60.0
            time_diff_3 = (row['get_time'] - row['put_time']).total_seconds() / 60.0

            # 使用传入的权重参数计算加权平均值
            return (self.weight_params['weight_1'] * time_diff_1) + \
                   (self.weight_params['weight_2'] * time_diff_2) + \
                   (self.weight_params['weight_3'] * time_diff_3)

        # 应用价值计算函数
        file_a['order_value'] = file_a.apply(calculate_order_value, axis=1)

        # 选择并重新排列所需的列
        transformed_data = file_a[
            ['start_timestamp', 'start_region', 'used', 'order_value', 'end_timestamp', 'end_region']]

        # 将订单数据结构化，存入订单列表
        # 订单数据结构： [(开始时间, 开始区域), 是否已被分配, 订单价值/奖励, (结束时间, 结束区域)]
        self.orders = transformed_data.apply(lambda row: [(int(row['start_timestamp']), int(row['start_region'])),
                                                          row['used'], row['order_value'],
                                                          (int(row['end_timestamp']), int(row['end_region']))],
                                             axis=1).tolist()
        return self.orders

    def get_orders_by_time(self, orders, t):
        """
        获取指定时间步的订单子集
        :param orders: 全部订单
        :param t: 时间步
        :return: 该时间步的订单子集
        """
        sub = []
        for i in range(0, len(self.orders)):
            if (orders[i][0][0] == t):
                sub.append(orders[i])
        return sub

    def learning(self):
        """
        使用时序差分法进行学习，更新价值矩阵
        """
        for time in range(self.time_count - 1, -1, -1):  # 从 T-1 到 0
            sub = self.get_orders_by_time(self.orders, time)  # 获取当前时间步的订单子集
            for time_sub in sub:
                R_t = 0  # 初始化奖励值
                self.N[time_sub[0][0]][time_sub[0][1]] += 1  # 更新区域的访问次数
                # 开始时间
                t1 = time_sub[0][0]
                # 结束时间
                t2 = time_sub[3][0]
                # 计算累计奖励 R_t
                for i in range(t2 - t1, -1, -1):
                    R_t += 2 * pow(0.9, i) * (time_sub[2]) / (t2 - t1 + 1)
                alpha = (1 / self.N[time_sub[0][0]][time_sub[0][1]])  # 学习率
                V_s_prime = self.V_s[time_sub[3][0]][time_sub[3][1]]  # 获取未来状态的价值
                # 更新当前状态的价值
                self.V_s[time_sub[0][0]][time_sub[0][1]] += alpha * (
                        (pow(0.9, t2 - t1 + 1)) * V_s_prime + R_t - self.V_s[time_sub[0][0]][time_sub[0][1]])

    def getV_s(self):
        """
        获取更新后的价值矩阵
        :return: 价值矩阵 V_s
        """
        return self.V_s

    def save_Vs(self, filepath='../data/Vs.npy'):
        """
        将 V_s 矩阵保存到文件中
        :param filepath: 保存文件的路径
        :param file_format: 文件格式，支持'npy'
        """
        np.save(filepath, self.V_s)
        print(f"V_s 矩阵已保存为 NumPy 的 .npy 文件：{filepath}")


if __name__ == '__main__':
    # 初始化 DatabaseReader
    db_reader = DatabaseReader()

    # 加权平均的权重参数，跟simulator中的权重保持一致
    weight_params = {'weight_1': 0.4, 'weight_2': 0.3, 'weight_3': 0.3}

    # code_index_map 文件路径
    code_index_map_filepath = '../data/node_index_with_cluster.xlsx'

    # 初始化 Learning 类，传入相关参数
    learning = Learning(
        time_count=144,  # 时间步数
        db_reader=db_reader,  # 数据库读取器实例
        code_index_map_filepath=code_index_map_filepath,  # code_index_map 文件路径
        weight_params=weight_params  # 加权参数
    )

    # 获取订单数据并执行学习算法
    data = learning.get_order_data()
    learning.learning()

    # 输出更新后的价值矩阵
    # print(learning.getV_s())

    # 保存价值矩阵
    learning.save_Vs()
    # 从文件中读取
    V_s_from_npy = np.load('../data/Vs.npy')
    print(V_s_from_npy.shape)
